// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <random>
#include <unordered_map>

#include <pollux/expression/function_signature.h>
#include <pollux/expression/signature_binder.h>
#include <pollux/type/type.h>
#include <pollux/vector/utils.h>

namespace kumo::pollux::fuzzer {

/// For function signatures using type variables, generates a list of
/// arguments types. Optionally, allows to specify a desired return type. If
/// specified, the return type acts as a constraint on the possible set of
/// argument types. If no return type is specified, it also allows generate a
/// random type that can bind to the function's return type.
class ArgumentTypeFuzzer {
 public:
  ArgumentTypeFuzzer(
      const exec::FunctionSignature& signature,
      FuzzerGenerator& rng)
      : ArgumentTypeFuzzer(signature, nullptr, rng) {}

  ArgumentTypeFuzzer(
      const exec::FunctionSignature& signature,
      const TypePtr& returnType,
      FuzzerGenerator& rng)
      : signature_{signature}, returnType_{returnType}, rng_{rng} {}

  /// Generate random argument types. If the desired returnType has been
  /// specified, checks that it can be bound to the return type of signature_.
  /// Return true if the generation succeeds, false otherwise. If signature_ has
  /// variable arity, repeat the last argument at most maxVariadicArgs times.
  bool fuzzArgumentTypes(uint32_t maxVariadicArgs);

  /// Return the generated list of argument types. This function should be
  /// called after fuzzArgumentTypes() is called and returns true.
  const std::vector<TypePtr>& argumentTypes() const {
    return argumentTypes_;
  }

  /// Return a random type that can bind to the function signature's return
  /// type and set returnType_ to this type. This function can only be called
  /// when returnType_ is uninitialized.
  TypePtr fuzzReturnType();

 private:
  /// Return the variables in the signature.
  auto& variables() const {
    return signature_.variables();
  }

  // Returns random integer between min and max inclusive.
  int32_t rand32(int32_t min, int32_t max);

  // Bind each type variable that is not determined by the return type to a
  // randomly generated type.
  void determineUnboundedTypeVariables();

  // Bind integer variables used in specified decimal(p,s) type signature to
  // randomly generated values if not already bound.
  // Noop if 'type' is not a decimal type signature.
  void determineUnboundedIntegerVariables(const exec::TypeSignature& type);

  TypePtr randType();

  /// Generates an orderable random type, including structs, and arrays.
  TypePtr randOrderableType();

  // Bind 'name' variable, if not already bound, using 'constant' constraint
  // ('name'='123'). Return bound value if 'name' is already bound or was
  // successfully bound to a constant value. Return std::nullopt otherwise.
  std::optional<int> tryFixedBinding(const std::string& name);

  // Bind the precision and scale variables in a decimal type signature to
  // constant values. Return std::nullopt if the variable cannot be bound to a
  // constant value.
  std::pair<std::optional<int>, std::optional<int>> tryBindFixedPrecisionScale(
      const exec::TypeSignature& type);

  // Find all the nested decimal type signatures recursively.
  void findDecimalTypes(
      const exec::TypeSignature& type,
      std::vector<exec::TypeSignature>& decimalTypes) const;

  const exec::FunctionSignature& signature_;

  TypePtr returnType_;

  std::vector<TypePtr> argumentTypes_;

  /// Bindings between type variables and their actual types.
  melon::F14FastMap<std::string, TypePtr> bindings_;

  melon::F14FastMap<std::string, int> integerBindings_;

  /// RNG to generate random types for unbounded type variables when necessary.
  FuzzerGenerator& rng_;
};

/// Return the kind name of type in lower case. This is expected to match the
/// TypeSignature::baseName_ corresponding to type.
std::string typeToBaseName(const TypePtr& type);

/// Return the TypeKind that corresponds to typeName.
std::optional<TypeKind> baseNameToTypeKind(const std::string& typeName);

} // namespace kumo::pollux::fuzzer
